# test_test_generator.py
import pytest
import asyncio
import json
import os
import tempfile
import uuid
from unittest.mock import Mock, patch, AsyncMock
from pathlib import Path
from datetime import datetime
from src.test_generator.TestGenerator import TestGenerator, ToolDefinition, TestCase


class TestTestGenerator:
    """TestGenerator 类的单元测试"""

    @pytest.fixture
    def sample_tool_definition(self):
        """提供示例工具定义"""
        return ToolDefinition(
            name="test_tool",
            description="A test tool for unit testing",
            input_schema={
                "type": "object",
                "properties": {
                    "param1": {"type": "string"},
                    "param2": {"type": "number"}
                }
            }
        )

    @pytest.fixture
    def sample_test_case(self):
        """提供示例测试用例"""
        return TestCase(
            id=str(uuid.uuid4()),
            toolName="test_tool",
            description="Test case description",
            query="Test query",
            input={"param1": "value1", "param2": 123},
            expect={
                "status": "success",
                "validation_rules": [
                    {"type": "contains", "value": "expected_string"}
                ]
            }
        )

    @pytest.fixture
    def mock_config(self):
        """模拟配置文件"""
        return {
            "mcpServers": {
                "test_server": {
                    "command": "python",
                    "args": ["-m", "test_server"],
                    "env": {"TEST_VAR": "test_value"}
                },
                "another_server": {
                    "command": "node",
                    "args": ["server.js"],
                    "env": {}
                }
            },
            "numTestsPerTool": 2
        }

    @pytest.fixture
    def temp_log_dir(self):
        """创建临时日志目录"""
        with tempfile.TemporaryDirectory() as temp_dir:
            yield temp_dir

    @pytest.fixture
    def test_generator(self, mock_config, temp_log_dir):
        """创建被测试的生成器实例"""
        with patch('src.client.MakeConfig.Configuration') as mock_config_class:
            mock_config_instance = Mock()
            mock_config_instance.load_config.return_value = mock_config
            mock_config_class.return_value = mock_config_instance
            
            with patch('src.llm.LLM.LLMClient') as mock_llm_class:
                mock_llm_instance = Mock()
                mock_llm_class.return_value = mock_llm_instance
                
                with patch('src.utils.read_source_code.ReadSourceCode') as mock_readsc:
                    mock_readsc_instance = Mock()
                    mock_readsc.return_value = mock_readsc_instance
                    mock_readsc_instance.get_code.return_value = {
                        "test_tool": "def test_tool(param1, param2):\n    return {'result': 'success'}"
                    }
                    
                    generator = TestGenerator(
                        api_key="test_api_key",
                        config_path="test/test_config.json",
                        log_name=temp_log_dir
                    )
                    
                    # 模拟 LLM 响应
                    mock_llm_instance.get_response.return_value = json.dumps([
                        {
                            "description": "Test case 1",
                            "input": {"param1": "value1"},
                            "expect": {
                                "status": "success",
                                "validationRules": [
                                    {"type": "contains", "value": "success"}
                                ]
                            }
                        }
                    ])
                    
                    yield generator

    @pytest.mark.asyncio
    async def test_initialization(self, test_generator):
        """测试初始化是否正确"""
        assert test_generator.log_name is not None
        assert test_generator.config is not None
        assert test_generator.config["numTestsPerTool"] == 2

    @pytest.mark.asyncio
    async def test_run_method_no_servers(self, test_generator):
        """测试没有服务器的情况"""
        test_generator.config["mcpServers"] = {}
        
        # 应该不会抛出异常
        await test_generator.run()

    @pytest.mark.asyncio
    async def test_run_method_with_servers(self, test_generator, mock_config):
        """测试运行方法（有服务器的情况）"""
        with patch('src.client.MCPClient.MCPClient') as mock_mcp_class:
            mock_mcp_instance = AsyncMock()
            mock_mcp_class.return_value = mock_mcp_instance
            
            # 模拟工具列表
            mock_tool = Mock()
            mock_tool.name = "test_tool"
            mock_tool.format_for_llm.return_value = "Tool: test_tool"
            mock_mcp_instance.list_tools.return_value = [mock_tool]
            mock_mcp_instance.name = "test_server"
            
            await test_generator.run()
            
            # 验证结果文件是否创建
            log_dir = Path(test_generator.log_name)
            test_dirs = [d for d in log_dir.iterdir() if d.is_dir() and d.name.startswith("test_server")]
            assert len(test_dirs) > 0
            
            # 验证测试用例文件
            testcase_file = test_dirs[0] / "testcases.json"
            assert testcase_file.exists()
            
            # 验证配置文件路径记录
            config_path_file = test_dirs[0] / "config_path.txt"
            assert config_path_file.exists()

    @pytest.mark.asyncio
    async def test_generate_tests_for_each_server_no_tools(self, test_generator):
        """测试生成测试用例（没有工具的情况）"""
        tools = []
        result = await test_generator.generate_tests_for_each_server(tools, 2, "test_server")
        
        assert result == []

    @pytest.mark.asyncio
    async def test_generate_tests_for_each_server_with_tools(self, test_generator, sample_tool_definition):
        """测试生成测试用例（有工具的情况）"""
        tools = [sample_tool_definition]
        
        result = await test_generator.generate_tests_for_each_server(tools, 2, "test_server")
        
        assert len(result) > 0
        assert isinstance(result[0], TestCase)
        assert result[0].toolName == "test_tool"

    @pytest.mark.asyncio
    async def test_generate_tests_for_each_server_llm_error(self, test_generator, sample_tool_definition):
        """测试生成测试用例（LLM 错误的情况）"""
        tools = [sample_tool_definition]
        
        # 模拟 LLM 错误
        test_generator.llm.get_response.side_effect = Exception("LLM API error")
        
        result = await test_generator.generate_tests_for_each_server(tools, 2, "test_server")
        
        assert result == []

    @pytest.mark.asyncio
    async def test_generate_tests_for_each_server_query_generation_error(self, test_generator, sample_tool_definition):
        """测试生成测试用例（查询生成错误的情况）"""
        tools = [sample_tool_definition]
        
        # 模拟查询生成错误
        def side_effect(messages):
            if "eval_prompt" in str(messages):
                raise Exception("Query generation failed")
            return json.dumps([{"description": "Test", "input": {}, "expect": {"status": "success"}}])
        
        test_generator.llm.get_response.side_effect = side_effect
        
        result = await test_generator.generate_tests_for_each_server(tools, 2, "test_server")
        
        # 应该仍然返回测试用例，但查询字段为空
        assert len(result) > 0
        assert result[0].query == ''

    def test_create_tool_prompt(self, test_generator, sample_tool_definition):
        """测试创建工具提示"""
        tool_function_str = "def test_tool(param):\n    return {'result': 'success'}"
        
        prompt = test_generator.create_tool_prompt(sample_tool_definition, 2, tool_function_str)
        
        assert prompt is not None
        assert "test_tool" in prompt
        assert "A test tool for unit testing" in prompt
        assert "2" in str(prompt)  # tests_per_tool

    def test_create_tool_prompt_no_input_schema(self, test_generator):
        """测试创建工具提示（没有输入模式的情况）"""
        tool = ToolDefinition(
            name="tool_no_schema",
            description="Tool without input schema",
            input_schema=None
        )
        
        prompt = test_generator.create_tool_prompt(tool, 2, "def tool(): pass")
        
        assert prompt is not None
        assert "tool_no_schema" in prompt

    def test_create_eval_prompt(self, test_generator, sample_tool_definition, sample_test_case):
        """测试创建评估提示"""
        prompt = test_generator.create_eval_prompt(sample_tool_definition, sample_test_case)
        
        assert prompt is not None
        assert "test_tool" in prompt
        assert "Test case description" in prompt

    def test_create_eval_prompt_no_input(self, test_generator, sample_tool_definition):
        """测试创建评估提示（没有输入的情况）"""
        test_case = TestCase(
            id=str(uuid.uuid4()),
            toolName="test_tool",
            description="Test case",
            query="",
            input=None,
            expect={"status": "success"}
        )
        
        prompt = test_generator.create_eval_prompt(sample_tool_definition, test_case)
        
        assert prompt is not None

    def test_parse_response_valid_json(self, test_generator):
        """测试解析响应（有效的 JSON）"""
        valid_response = json.dumps([
            {
                "description": "Valid test case",
                "input": {"param": "value"},
                "expect": {
                    "status": "success",
                    "validationRules": [{"type": "contains", "value": "success"}]
                }
            }
        ])
        
        result = test_generator.parse_response(valid_response, "test_tool")
        
        assert len(result) == 1
        assert result[0].toolName == "test_tool"
        assert result[0].description == "Valid test case"

    def test_parse_response_with_backticks(self, test_generator):
        """测试解析响应（带反引号的 JSON）"""
        response_with_backticks = f"""
                    Here are the test cases:```json
            {json.dumps([
            {
            "description": "Test case with backticks",
            "input": {"param": "value"},
            "expect": {
            "status": "success",
            "validationRules": []
            }
            }
            ])}
        """

        result = test_generator.parse_response(response_with_backticks, "test_tool")

        assert len(result) == 1
        assert result[0].description == "Test case with backticks"

    def test_parse_response_invalid_json(self, test_generator):
        """测试解析响应（无效的 JSON）"""
        invalid_response = "This is not valid JSON"

        result = test_generator.parse_response(invalid_response, "test_tool")

        assert result == []

    def test_parse_response_missing_required_fields(self, test_generator):
        """测试解析响应（缺少必需字段）"""
        incomplete_response = json.dumps([
            {
                "description": "Test case",
                # 缺少 input 字段
                "expect": {
                    "status": "success"
                }
            }
        ])

        result = test_generator.parse_response(incomplete_response, "test_tool")

        assert result == []

    def test_parse_response_invalid_status(self, test_generator):
        """测试解析响应（无效的状态）"""
        response_with_invalid_status = json.dumps([
            {
                "description": "Test case",
                "input": {"param": "value"},
                "expect": {
                    "status": "invalid_status",  # 无效状态
                    "validationRules": []
                }
            }
        ])

        result = test_generator.parse_response(response_with_invalid_status, "test_tool")

        assert result == []

    def test_parse_response_single_object(self, test_generator):
        """测试解析响应（单个对象而非数组）"""
        single_object_response = json.dumps({
            "description": "Single test case",
            "input": {"param": "value"},
            "expect": {
                "status": "success",
                "validationRules": []
            }
        })

        result = test_generator.parse_response(single_object_response, "test_tool")

        assert len(result) == 1  # 应该被包装成数组
        assert result[0].description == "Single test case"

    def test_parse_response_empty_validation_rules(self, test_generator):
        """测试解析响应（空的验证规则）"""
        response_with_empty_rules = json.dumps([
            {
                "description": "Test case",
                "input": {"param": "value"},
                "expect": {
                    "status": "success",
                    "validationRules": []  # 空规则列表
                }
            }
        ])

        result = test_generator.parse_response(response_with_empty_rules, "test_tool")

        assert len(result) == 1
        assert result[0].expect["validation_rules"] == []

    def test_parse_response_unexpected_error(self, test_generator):
        """测试解析响应（意外错误）"""
        # 模拟解析过程中的意外错误
        with patch('json.loads', side_effect=Exception("Unexpected error")):
            result = test_generator.parse_response("{}", "test_tool")
            assert result == []

    def test_testcases_to_dict(self, test_generator, sample_test_case):
        """测试测试用例转换为字典"""
        test_cases = [sample_test_case]

        result = test_generator.testcases_to_dict(test_cases)

        assert len(result) == 1
        assert result[0]["id"] == sample_test_case.id
        assert result[0]["toolName"] == sample_test_case.toolName
        assert result[0]["description"] == sample_test_case.description
        assert result[0]["query"] == sample_test_case.query
        assert result[0]["input"] == sample_test_case.input
        assert result[0]["expect"] == sample_test_case.expect

    def test_testcases_to_dict_empty_list(self, test_generator):
        """测试测试用例转换为字典（空列表）"""
        result = test_generator.testcases_to_dict([])

        assert result == []

    def test_save_to_file_success(self, test_generator, sample_test_case, temp_log_dir):
        """测试保存到文件成功的情况"""
        test_cases = [sample_test_case]

        result = test_generator.save_to_file("test_server", test_cases)

        assert result is True

        # 验证文件是否创建
        log_dir = Path(temp_log_dir)
        test_dirs = [d for d in log_dir.iterdir() if d.is_dir() and d.name.startswith("test_server")]
        assert len(test_dirs) > 0

        testcase_file = test_dirs[0] / "testcases.json"
        assert testcase_file.exists()

        # 验证文件内容
        with open(testcase_file, 'r', encoding='utf-8') as f:
            content = json.load(f)
            assert len(content) == 1
            assert content[0]["toolName"] == "test_tool"

    def test_save_to_file_invalid_data(self, test_generator):
        """测试保存到文件（无效数据的情况）"""
        # 传递非列表数据应该返回 False
        result = test_generator.save_to_file("test_server", "invalid_data")
        assert result is False

    def test_save_to_file_io_error(self, test_generator, sample_test_case):
        """测试保存到文件（IO 错误的情况）"""
        # 模拟 IO 错误
        with patch('builtins.open', side_effect=IOError("Disk full")):
            result = test_generator.save_to_file("test_server", [sample_test_case])
            assert result is False

    def test_save_to_file_directory_creation_error(self, test_generator, sample_test_case):
        """测试保存到文件（目录创建错误的情况）"""
        # 模拟目录创建错误
        with patch('os.mkdir', side_effect=OSError("Permission denied")):
            result = test_generator.save_to_file("test_server", [sample_test_case])
            assert result is False

    @pytest.mark.asyncio
    async def test_tool_initialization_error(self, test_generator, mock_config):
        """测试工具初始化错误的情况"""
        with patch('src.client.MCPClient.MCPClient') as mock_mcp_class:
            mock_mcp_instance = AsyncMock()
            mock_mcp_class.return_value = mock_mcp_instance
            
            # 模拟初始化错误
            mock_mcp_instance.initialize.side_effect = Exception("Connection failed")
            mock_mcp_instance.name = "test_server"
            
            # 应该不会抛出异常
            await test_generator.run()

    @pytest.mark.asyncio
    async def test_tool_list_empty(self, test_generator, mock_config):
        """测试工具列表为空的情况"""
        with patch('src.client.MCPClient.MCPClient') as mock_mcp_class:
            mock_mcp_instance = AsyncMock()
            mock_mcp_class.return_value = mock_mcp_instance
            
            # 模拟空工具列表
            mock_mcp_instance.list_tools.return_value = []
            mock_mcp_instance.name = "test_server"
            
            await test_generator.run()
            
            # 验证没有测试用例文件创建
            log_dir = Path(test_generator.log_name)
            test_dirs = [d for d in log_dir.iterdir() if d.is_dir() and d.name.startswith("test_server")]
            assert len(test_dirs) == 0

    @pytest.mark.asyncio
    async def test_multiple_servers(self, test_generator, mock_config):
        """测试多个服务器的情况"""
        with patch('src.client.MCPClient.MCPClient') as mock_mcp_class:
            # 模拟第一个服务器
            mock_server1 = AsyncMock()
            mock_server1.name = "test_server"
            mock_server1.list_tools.return_value = []
            
            # 模拟第二个服务器
            mock_server2 = AsyncMock()
            mock_server2.name = "another_server"
            mock_server2.list_tools.return_value = []
            
            # 让 MCPClient 依次返回不同的服务器实例
            mock_mcp_class.side_effect = [mock_server1, mock_server2]
            
            await test_generator.run()
            
            # 验证两个服务器的目录都创建了
            log_dir = Path(test_generator.log_name)
            test_dirs = [d for d in log_dir.iterdir() if d.is_dir()]
            server_names = [d.name.split('_')[0] for d in test_dirs]
            
            assert "test_server" in server_names
            assert "another_server" in server_names

    def test_timestamp_formatting(self, test_generator, sample_test_case):
        """测试时间戳格式化"""
        # 模拟一个特定的时间戳
        with patch('datetime.datetime') as mock_datetime:
            fixed_time = datetime(2023, 10, 15, 12, 30, 45, 123456)
            mock_datetime.utcnow.return_value = fixed_time
            
            result = test_generator.save_to_file("test_server", [sample_test_case])
            
            # 验证目录名称格式
            log_dir = Path(test_generator.log_name)
            test_dirs = [d for d in log_dir.iterdir() if d.is_dir()]
            
            # 应该包含格式化的时间戳
            assert any("2023-10-15T12-30-45-123456" in str(d) for d in test_dirs)


    # 测试 ToolDefinition 和 TestCase 类
    class TestToolDefinition:
        """ToolDefinition 类的测试"""

        def test_tool_definition_creation(self):
            """测试工具定义创建"""
            tool = ToolDefinition(
                name="test_tool",
                description="Test tool",
                input_schema={"type": "object"}
            )

            assert tool.name == "test_tool"
            assert tool.description == "Test tool"
            assert tool.input_schema == {"type": "object"}


    class TestTestCase:
        """TestCase 类的测试"""
        def test_test_case_creation(self):
            """测试测试用例创建"""
            test_id = str(uuid.uuid4())
            test_case = TestCase(
                id=test_id,
                toolName="test_tool",
                description="Test case",
                query="Test query",
                input={"param": "value"},
                expect={"status": "success"}
            )

            assert test_case.id == test_id
            assert test_case.toolName == "test_tool"
            assert test_case.description == "Test case"
            assert test_case.query == "Test query"
            assert test_case.input == {"param": "value"}
            assert test_case.expect == {"status": "success"}


if __name__ == "__main__":
    pytest.main([__file__, "-v"])